import glm
import OpenGL.GL as gl
import numpy as np
from pygl import *

class GridHelper:
    def __init__(self, size=2, divisions=20, color1=glm.vec3(0.0, 1.0, 0.0), 
                 color2=glm.vec3(1.0, 1.0, 1.0)) -> None:
        self.program = ProgramVF("common/shaders/axes.vs",
                                "common/shaders/axes.fs")
        center = divisions / 2
        step = size / divisions
        half_size = size / 2
        vertices = []
        colors = []
        k = -half_size
        for i in range(divisions + 1):
            vertices.append([-half_size, 0.0, k])
            vertices.append([half_size, 0.0, k])
            vertices.append([k, 0.0, -half_size])
            vertices.append([k, 0.0, half_size])
            if i == center:
                colors.append([color1.x, color1.y, color1.z])
                colors.append([color1.x, color1.y, color1.z])
                colors.append([color1.x, color1.y, color1.z])
                colors.append([color1.x, color1.y, color1.z])
            else:    
                colors.append([color2.x, color2.y, color2.z])
                colors.append([color2.x, color2.y, color2.z])
                colors.append([color2.x, color2.y, color2.z])
                colors.append([color2.x, color2.y, color2.z])
            k += step
        self.num_vertices = len(vertices)
        self.position_vbo = VertexBufferObject(np.array(vertices, 
                                            dtype=gl.GLfloat))
        self.color_vbo = VertexBufferObject(np.array(colors, dtype=gl.GLfloat))
        self.vao = VertexArrayObject()
        position_binding_point = 0
        color_binding_point = 1
        self.vao.setVertexBuffer(self.position_vbo, position_binding_point, 0, 
                                 3 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexBuffer(self.color_vbo, color_binding_point, 0, 
                                 3 * gl.sizeof(gl.GLfloat))
        attribute_position = VertexAttribute("position", 0, 3, gl.GL_FLOAT, 
                                             gl.GL_FALSE, 0)
        attribute_color = VertexAttribute("color",1, 3, gl.GL_FLOAT,
                                           gl.GL_FALSE, 0)
        self.vao.setVertexAttribute(position_binding_point, attribute_position)
        self.vao.setVertexAttribute(color_binding_point, attribute_color)

    def render(self, view: glm.mat4, projection: glm.mat4) -> None:
        self.program.use()
        model = glm.mat4(1.0)
        self.program.setUniformMatrix4fv(0, model)
        self.program.setUniformMatrix4fv(1, view)
        self.program.setUniformMatrix4fv(2, projection)
        gl.glLineWidth(1.0)
        self.vao.bind()
        gl.glDrawArrays(gl.GL_LINES, 0, self.num_vertices)

    def delete(self)->None:
        self.vao.delete()
        self.color_vbo.delete()
        self.position_vbo.delete()
        self.program.delete()

